{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKwBhBWZlhOV"
      },
      "source": [
        "Copyright 2024 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use\n",
        "this file except in compliance with the License. You may obtain a copy of the\n",
        "License at\n",
        "\n",
        "```\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "```\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed\n",
        "under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR\n",
        "CONDITIONS OF ANY KIND, either express or implied. See the License for the\n",
        "specific language governing permissions and limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e-_uYcH8vCvF"
      },
      "source": [
        "# Simulation study\n",
        "\n",
        "This code constructs a simulated data generating process where there are three\n",
        "proxies $Y_1$, $Y_2$, and $Y_3$ of a variable $Y$. Dependent on the\n",
        "data-generating process, $Y$ is generated by covariates $X$ and (depending on\n",
        "the data-generating process) subgroup membership $A$. $Y_1$ and $Y_2$ are\n",
        "\"unbiased proxies\" of $Y$ because they depend only on $Y$. $Y_2$ is noisier than\n",
        "$Y_1$. $Y_3$ is a biased proxy because it depends on both $Y$ and $A$.\n",
        "\n",
        "We fit models to predict $Y$ and each of its proxies using $X$ or $\\{X, A\\}$ for\n",
        "a collection of data generating processes. We evaluate the sufficiency fairness\n",
        "criterion ($Y_i \\perp A \\mid R$ for a continuous score $R$ and label $Y_i$) using calibration curves."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijf85jeulwRW"
      },
      "outputs": [],
      "source": [
        "from causal_label_bias import utils\n",
        "import jax\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from scipy.special import expit\n",
        "import seaborn as sns\n",
        "from sklearn.preprocessing import OneHotEncoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p8Xl9blluIWy"
      },
      "outputs": [],
      "source": [
        "class Simulator:\n",
        "  \"\"\"Generates simulated data following a causal generative process\"\"\"\n",
        "\n",
        "  def __init__(self, **kwargs):\n",
        "    \"\"\"Initializes simulation.\n",
        "\n",
        "    Arguments:\n",
        "      param_dict: A dictionary for data-generating parameters. The provided\n",
        "        param_dict will override default parameters.\n",
        "\n",
        "    Returns:\n",
        "      A pd.DataFrame with columns corresponding to the variables in the causal\n",
        "      DAG.\n",
        "    \"\"\"\n",
        "\n",
        "    self.param_dict = self.get_default_param_dict()\n",
        "\n",
        "    if kwargs is not None:\n",
        "      for key, value in kwargs.items():\n",
        "        self.param_dict[key] = value\n",
        "\n",
        "  def get_default_param_dict(self):\n",
        "    param_dict = {\n",
        "        'num_samples': 5000,\n",
        "        'k_x': 1,\n",
        "        'k_y': 1,\n",
        "        'mu_x_u': np.array([0.0, 0.0]),\n",
        "        'beta_a': 0,\n",
        "        'pi_a': np.array([0.5, 0.5]),\n",
        "        'mu_y1_a': np.array([0.0, 0.0]),\n",
        "        'mu_y2_a': np.array([0.0, 0.0]),\n",
        "        'mu_y3_a': np.array([0.0, 0.5]),\n",
        "        'mu_y_a': np.array([0, 2]),\n",
        "        'mu_y_x': 1,\n",
        "        'mu_y1_y': np.array([1]),\n",
        "        'mu_y2_y': np.array([0.5]),\n",
        "        'mu_y3_y': np.array([1.5]),\n",
        "        'mu_y_x_base': 2,\n",
        "        'sd_x': 3,\n",
        "        'p_u': [0.5, 0.5],\n",
        "        'loc_y1': 0.0,\n",
        "        'scale_y1': 0.5,\n",
        "        'loc_y2': 0.0,\n",
        "        'scale_y2': 0.5,\n",
        "        'loc_y3': 0.0,\n",
        "        'scale_y3': 0.0,\n",
        "    }\n",
        "    return param_dict\n",
        "\n",
        "  def get_samples(self, p_u=None, seed=42):\n",
        "    \"\"\"Generates samples from the simulation.\n",
        "\n",
        "    Arguments:\n",
        "      p_u: array that specifies the mixture proportions over latent categories u\n",
        "      seed: a random seed\n",
        "\n",
        "    Returns:\n",
        "      a dict containing generated data\n",
        "    \"\"\"\n",
        "\n",
        "    rng = jax.random.PRNGKey(seed)\n",
        "    _, k0, _ = jax.random.split(rng, 3)\n",
        "\n",
        "    ## Generate u\n",
        "    if p_u is None:\n",
        "      p_u = self.param_dict['p_u']\n",
        "\n",
        "    u = np.random.binomial(1, p_u[1], size=self.param_dict['num_samples'])\n",
        "    u_one_hot = OneHotEncoder(sparse=False).fit_transform(u.reshape(-1, 1))\n",
        "\n",
        "    ## Generate x\n",
        "    x = jax.random.multivariate_normal(\n",
        "        key=k0,\n",
        "        mean=(u_one_hot @ self.param_dict['mu_x_u']).reshape(-1, 1),\n",
        "        cov=self.param_dict['sd_x'] * np.eye(self.param_dict['k_x']),\n",
        "    )\n",
        "    x = np.array(x).astype(np.float64)\n",
        "\n",
        "    ## Generate a\n",
        "    p_a = (\n",
        "        self.param_dict['beta_a'] * u_one_hot\n",
        "        + (1 - self.param_dict['beta_a']) * self.param_dict['pi_a']\n",
        "    )\n",
        "    a = np.random.binomial(1, p_a[:, 1], size=self.param_dict['num_samples'])\n",
        "    a_one_hot = OneHotEncoder(sparse=False).fit_transform(a.reshape(-1, 1))\n",
        "\n",
        "    ## Generate y\n",
        "    mu_y_x = np.array(\n",
        "        [[self.param_dict['mu_y_x'], -1 * self.param_dict['mu_y_x']]]\n",
        "    )\n",
        "\n",
        "    y = x.dot(mu_y_x)[\n",
        "        np.arange(self.param_dict['num_samples']), np.squeeze(a)\n",
        "    ].reshape(-1, 1) + (a_one_hot @ self.param_dict['mu_y_a']).reshape(-1, 1)\n",
        "    p_y = expit(y)\n",
        "    y_bin = np.squeeze(np.random.binomial(n=1, p=p_y))\n",
        "\n",
        "    ## Generate y1\n",
        "\n",
        "    y1_logits = y.dot(self.param_dict['mu_y1_y']).reshape(\n",
        "        -1, 1\n",
        "    ) + np.random.normal(\n",
        "        loc=self.param_dict['loc_y1'],\n",
        "        scale=self.param_dict['scale_y1'],\n",
        "        size=self.param_dict['num_samples'],\n",
        "    ).reshape(\n",
        "        -1, 1\n",
        "    )\n",
        "\n",
        "    p_y1 = expit(y1_logits)\n",
        "    y1 = np.squeeze(np.random.binomial(n=1, p=p_y1))\n",
        "    y1_one_hot = OneHotEncoder(sparse=False).fit_transform(y1.reshape(-1, 1))\n",
        "\n",
        "    ## Generate y2\n",
        "    y2_logits = y.dot(self.param_dict['mu_y2_y']).reshape(\n",
        "        -1, 1\n",
        "    ) + np.random.normal(\n",
        "        loc=self.param_dict['loc_y2'],\n",
        "        scale=self.param_dict['scale_y2'],\n",
        "        size=self.param_dict['num_samples'],\n",
        "    ).reshape(\n",
        "        -1, 1\n",
        "    )\n",
        "\n",
        "    p_y2 = expit(y2_logits)\n",
        "    y2 = np.squeeze(np.random.binomial(n=1, p=p_y2))\n",
        "    y2_one_hot = OneHotEncoder(sparse=False).fit_transform(y2.reshape(-1, 1))\n",
        "\n",
        "    ## Generate y3\n",
        "\n",
        "    y3_logits = (\n",
        "        y.dot(self.param_dict['mu_y3_y']).reshape(-1, 1)\n",
        "        + (np.tan(a_one_hot @ self.param_dict['mu_y3_a']).reshape(-1, 1))\n",
        "        + np.random.normal(\n",
        "            loc=self.param_dict['loc_y3'],\n",
        "            scale=self.param_dict['scale_y3'],\n",
        "            size=self.param_dict['num_samples'],\n",
        "        ).reshape(-1, 1)\n",
        "    )\n",
        "\n",
        "    p_y3 = expit(y3_logits)\n",
        "    y3 = np.squeeze(np.random.binomial(n=1, p=p_y3))\n",
        "    y3_one_hot = OneHotEncoder(sparse=False).fit_transform(y3.reshape(-1, 1))\n",
        "\n",
        "    return {\n",
        "        'u': u,\n",
        "        'a': a,\n",
        "        'x': x,\n",
        "        'y': y,\n",
        "        'y_bin': y_bin,\n",
        "        'p_y': p_y,\n",
        "        'y1': y1,\n",
        "        'y1_logits': y1_logits,\n",
        "        'p_y1': p_y1,\n",
        "        'y1_one_hot': y1_one_hot,\n",
        "        'y2': y2,\n",
        "        'y2_logits': y2_logits,\n",
        "        'p_y2': p_y2,\n",
        "        'y2_one_hot': y2_one_hot,\n",
        "        'y3': y3,\n",
        "        'y3_logits': y3_logits,\n",
        "        'p_y3': p_y3,\n",
        "        'y3_one_hot': y3_one_hot,\n",
        "    }\n",
        "\n",
        "\n",
        "def get_squeezed_df(data_dict: dict) -\u003e pd.DataFrame:\n",
        "  \"\"\"Converts a dict of numpy arrays into a DataFrame, extracting columns of arrays into separate DataFrame columns.\"\"\"\n",
        "  temp = {}\n",
        "  for key, value in data_dict.items():\n",
        "    squeezed_array = np.squeeze(value)\n",
        "    if len(squeezed_array.shape) == 1:\n",
        "      temp[key] = squeezed_array\n",
        "    elif len(squeezed_array.shape) \u003e 1:\n",
        "      for i in range(value.shape[1]):\n",
        "        temp[f'{key}_{i}'] = np.squeeze(value[:, i])\n",
        "  df = pd.DataFrame(temp)\n",
        "  return df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dxs2T9ineL7n"
      },
      "outputs": [],
      "source": [
        "def plot_calibration_curves(\n",
        "    the_dict: dict,\n",
        "    x_var: str = 'pred_probs',\n",
        "    y_var: str = 'calibration_curve',\n",
        "    ax=None,\n",
        "    show_legend: bool = False,\n",
        "    **kwargs,\n",
        "):\n",
        "  \"\"\"Plot calibration curves.\n",
        "\n",
        "  Arguments:\n",
        "    the_dict: A dictionary containing values to be plotted for each subgroup,\n",
        "      where subgroup is the key.\n",
        "    x_var: A string specifying the x-axis variable.\n",
        "    y_var: A string specifying the y-axis variable.\n",
        "  \"\"\"\n",
        "  if ax is None:\n",
        "    plt.figure()\n",
        "    ax = plt.gca()\n",
        "  for key in sorted(the_dict.keys()):\n",
        "    ax.plot(the_dict[key][x_var], the_dict[key][y_var], label=key, **kwargs)\n",
        "\n",
        "  ax.plot(\n",
        "      np.linspace(0, 1, 100),\n",
        "      np.linspace(0, 1, 100),\n",
        "      alpha=0.5,\n",
        "      linestyle='--',\n",
        "      color='k',\n",
        "      label='Identity',\n",
        "  )\n",
        "  sns.despine(ax=ax)\n",
        "  if show_legend:\n",
        "    plt.legend(sorted(the_dict.keys()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "isG_-9KDkkEk"
      },
      "outputs": [],
      "source": [
        "def plot(evals: dict, title: str):\n",
        "  \"\"\"Plot calibration curves for different setups\n",
        "\n",
        "  Arguments:\n",
        "    evals: A dict containing dicts for different outcomes. Each inner dict is\n",
        "      keyed by group, contains dataframes containing the fitted models. The\n",
        "      dataframe is keyed by \"outcomes\" (true outcomes) and \"pred_probs\"\n",
        "      (predicted probabilities generated by the fitted model).\n",
        "    title: A string containing the title for the plot\n",
        "  \"\"\"\n",
        "  plt.close()\n",
        "  plt.rcParams['xtick.labelsize'] = 12\n",
        "  plt.rcParams['ytick.labelsize'] = 12\n",
        "  n_eval_outcomes = len(evals['y'])\n",
        "  _, ax = plt.subplots(\n",
        "      4, n_eval_outcomes, figsize=(18, 12), sharey=True, sharex=True\n",
        "  )\n",
        "  plt.subplots_adjust(wspace=0.1, hspace=0.05)\n",
        "\n",
        "  ax[0][0].set_ylabel(r'P(Outcome)$', fontsize=24)\n",
        "  ax[1][0].set_ylabel(r'P(Outcome)$', fontsize=24)\n",
        "  ax[2][0].set_ylabel(r'P(Outcome)$', fontsize=24)\n",
        "  ax[3][0].set_ylabel(r'P(Outcome)$', fontsize=24)\n",
        "\n",
        "  ax[0][0].set_title('Eval Outcome: $Y_1$ ', fontsize=24)\n",
        "  ax[0][1].set_title('Eval Outcome: $Y_2$', fontsize=24)\n",
        "  ax[0][2].set_title('Eval Outcome: $Y_3$', fontsize=24)\n",
        "  ax[0][3].set_title('Eval Outcome: $Y$', fontsize=24)\n",
        "\n",
        "  # Plot calibration curve with proxy outcome y1\n",
        "  for i in range(len(evals['y1'])):\n",
        "    plot_calibration_curves(\n",
        "        the_dict=evals['y1'][i],\n",
        "        x_var='pred_probs',\n",
        "        y_var='calibration_curve_y',\n",
        "        ax=ax[0][i],\n",
        "        lw=3,\n",
        "    )\n",
        "\n",
        "  # Plot calibration curve with proxy outcome y2\n",
        "  for i in range(len(evals['y2'])):\n",
        "    plot_calibration_curves(\n",
        "        the_dict=evals['y2'][i],\n",
        "        x_var='pred_probs',\n",
        "        y_var='calibration_curve_y',\n",
        "        ax=ax[1][i],\n",
        "        lw=3,\n",
        "    )\n",
        "\n",
        "  # Plot calibration curve with proxy outcome y2\n",
        "  for i in range(len(evals['y3'])):\n",
        "    plot_calibration_curves(\n",
        "        the_dict=evals['y3'][i],\n",
        "        x_var='pred_probs',\n",
        "        y_var='calibration_curve_y',\n",
        "        ax=ax[2][i],\n",
        "        lw=3,\n",
        "    )\n",
        "\n",
        "  # Plot calibration curve with true outcome\n",
        "  for i in range(len(evals['y'])):\n",
        "    plot_calibration_curves(\n",
        "        the_dict=evals['y'][i],\n",
        "        x_var='pred_probs',\n",
        "        y_var='calibration_curve_y',\n",
        "        ax=ax[3][i],\n",
        "        lw=3,\n",
        "    )\n",
        "\n",
        "  plt.gcf().text(x=1, y=0.80, s='Train Outcome: $Y_1$', fontsize=24)\n",
        "  plt.gcf().text(x=1, y=0.60, s='Train Outcome: $Y_2$', fontsize=24)\n",
        "  plt.gcf().text(x=1, y=0.40, s='Train Outcome: $Y_3$', fontsize=24)\n",
        "  plt.gcf().text(x=1, y=0.20, s='Train Outcome: $Y$', fontsize=24)\n",
        "  plt.gcf().text(x=0.4, y=0.05, s='Predicted probability', fontsize=24)\n",
        "  plt.gcf().text(x=0.35, y=1, s='Causal graph has ' + title, fontsize=34)\n",
        "\n",
        "  sns.despine()\n",
        "  plt.figure()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HcOcqAmheIFO"
      },
      "outputs": [],
      "source": [
        "def train_setups(\n",
        "    dgp_type: str = 'no_a',\n",
        "    eval_outcomes: str = ['y1', 'y2', 'y3', 'y_bin'],\n",
        "    features: str = ['x'],\n",
        "    sim_samples_dict: dict = {},\n",
        "    sim_samples_dict_eval: dict = {},\n",
        "    model_type: str = 'gradient_boosting',\n",
        "    stratified: bool = False,\n",
        "):\n",
        "  \"\"\"Train and evaluate models with different outcomes.\n",
        "\n",
        "  Arguments:\n",
        "    dgp_type: A string specifying edge connections in the causal DAG.\n",
        "    eval_outcomes: An iterable of strings indicating the different evaluation\n",
        "      outcomes.\n",
        "    features: A string or iterable of strings indicating the name(s) of the\n",
        "      columns used as features.\n",
        "    sim_samples_dict: A dictionary containing a pd.DataFrame for model training.\n",
        "    sim_samples_dict_eval: A dictionary containing a pd.DataFrame for model\n",
        "      evaluation.\n",
        "    model_type: A string specifying the model type.\n",
        "    stratified: A boolean flag for training group specific models.\n",
        "\n",
        "  Returns:\n",
        "    An a dict containing dicts for different training outcomes.\n",
        "    Each dict keyed by group with dict values.\n",
        "    The inner dicts contain, for each group, the fitted model and numpy arrays\n",
        "    corresponding to the features, labels, and predictions.\n",
        "  \"\"\"\n",
        "\n",
        "  evals = dict()\n",
        "  # outcome for training is y1\n",
        "  evals['y1'] = [\n",
        "      utils.fit_models_df(\n",
        "          source_df=sim_samples_dict[dgp_type],\n",
        "          target_df=sim_samples_dict_eval[dgp_type],\n",
        "          outcome_key='y1',\n",
        "          outcome_key_target=eval_outcome,\n",
        "          features_keys=features,\n",
        "          model_type=model_type,\n",
        "          stratified=stratified,\n",
        "      )\n",
        "      for eval_outcome in eval_outcomes\n",
        "  ]\n",
        "  # outcome for training is y2\n",
        "  evals['y2'] = [\n",
        "      utils.fit_models_df(\n",
        "          source_df=sim_samples_dict[dgp_type],\n",
        "          target_df=sim_samples_dict_eval[dgp_type],\n",
        "          outcome_key='y2',\n",
        "          outcome_key_target=eval_outcome,\n",
        "          features_keys=features,\n",
        "          model_type=model_type,\n",
        "          stratified=stratified,\n",
        "      )\n",
        "      for eval_outcome in eval_outcomes\n",
        "  ]\n",
        "  # outcome for training is y3\n",
        "  evals['y3'] = [\n",
        "      utils.fit_models_df(\n",
        "          source_df=sim_samples_dict[dgp_type],\n",
        "          target_df=sim_samples_dict_eval[dgp_type],\n",
        "          outcome_key='y3',\n",
        "          outcome_key_target=eval_outcome,\n",
        "          features_keys=features,\n",
        "          model_type=model_type,\n",
        "          stratified=stratified,\n",
        "      )\n",
        "      for eval_outcome in eval_outcomes\n",
        "  ]\n",
        "  # outcome for training is h_bin\n",
        "  evals['y'] = [\n",
        "      utils.fit_models_df(\n",
        "          source_df=sim_samples_dict[dgp_type],\n",
        "          target_df=sim_samples_dict_eval[dgp_type],\n",
        "          outcome_key='y_bin',\n",
        "          outcome_key_target=eval_outcome,\n",
        "          features_keys=features,\n",
        "          model_type=model_type,\n",
        "          stratified=stratified,\n",
        "      )\n",
        "      for eval_outcome in eval_outcomes\n",
        "  ]\n",
        "  return evals"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dGO0PPqyK9Kb"
      },
      "outputs": [],
      "source": [
        "sim_dict = {\n",
        "    'y3_a_no_y_a': Simulator(\n",
        "        mu_y3_a=np.array([0, -2]), mu_y_a=np.array([0, 0]), scale_y2=15\n",
        "    ),\n",
        "    'y_a_no_y3_a': Simulator(\n",
        "        mu_y3_a=np.array([0, 0]), mu_y_a=np.array([-1, 1]), scale_y2=15\n",
        "    ),\n",
        "    'no_a': Simulator(\n",
        "        mu_y3_a=np.array([0, 0]), mu_y_a=np.array([0, 0]), scale_y2=15\n",
        "    ),\n",
        "    'y_a_y3_a': Simulator(\n",
        "        mu_y3_a=np.array([0, -2]), mu_y_a=np.array([-1, 1]), scale_y2=4\n",
        "    ),\n",
        "}\n",
        "sim_samples_dict = {\n",
        "    key: get_squeezed_df(value.get_samples(seed=i))\n",
        "    for i, (key, value) in enumerate(sim_dict.items())\n",
        "}\n",
        "sim_samples_dict_eval = {\n",
        "    key: get_squeezed_df(value.get_samples(seed=2 * i))\n",
        "    for i, (key, value) in enumerate(sim_dict.items())\n",
        "}\n",
        "DISPLAY_PLOTS = True\n",
        "if DISPLAY_PLOTS:\n",
        "  plt.close()\n",
        "  for key, value in sim_samples_dict.items():\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='x', hue='a')\n",
        "    plt.title(f'{key}: x')\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='y', hue='a')\n",
        "    plt.title(f'{key}: y')\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='p_y', hue='a')\n",
        "    plt.title(f'{key}: p_y')\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='y1_logits', hue='a')\n",
        "    plt.title(f'{key}: y1_logits')\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='y2_logits', hue='a')\n",
        "    plt.title(f'{key}: y2_logits')\n",
        "    plt.figure()\n",
        "    sns.kdeplot(value, x='y3_logits', hue='a')\n",
        "    plt.title(f'{key}: y3_logits')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w2Ctj_eglZDL"
      },
      "source": [
        "# Scenario 1: Model trained with [$X$]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ds55z4J8UD1_"
      },
      "outputs": [],
      "source": [
        "dgp_types = ['no_a', 'y3_a_no_y_a', 'y_a_no_y3_a', 'y_a_y3_a']\n",
        "eval_outcomes = ['y1', 'y2', 'y3', 'y_bin']\n",
        "features = ['x']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SlJUq9IOFWQ6"
      },
      "source": [
        "### This refers to the data generating process without $A → Y$ or $A → Y_3$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TXU47o0eV_D2"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[0]))\n",
        "evals = train_setups(\n",
        "    dgp_types[0],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "avyWjcYUFtQE"
      },
      "source": [
        "### This refers to the data generating process without $A \\rightarrow Y$, but with $A → Y_3$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w4njvz_CWV7Y"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[1]))\n",
        "evals = train_setups(\n",
        "    dgp_types[1],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOV3f9LeGOK6"
      },
      "source": [
        "### This refers to the data generating process with $A → Y$, but without $A → Y_3$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KcnWSVOVWXkV"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[2]))\n",
        "evals = train_setups(\n",
        "    dgp_types[2],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5mJLmWcHGZuR"
      },
      "source": [
        "### This refers to the data generating process with $A → Y$ and $A → Y_3$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "awDT2PgBWZ8d"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[3]))\n",
        "evals = train_setups(\n",
        "    dgp_types[3],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[3])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nDpbK9a2mAgM"
      },
      "source": [
        "# Scenario 2: Models trained with [$X$,$A$]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vNZD1yZ-mfVl"
      },
      "outputs": [],
      "source": [
        "features = ['x', 'a']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zA14RYdgW6bQ"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[0]))\n",
        "evals = train_setups(\n",
        "    dgp_types[0],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xTHCDdYFSro1"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[1]))\n",
        "evals = train_setups(\n",
        "    dgp_types[1],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h8TCuBaoXOQy"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[2]))\n",
        "evals = train_setups(\n",
        "    dgp_types[2],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[2])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EHu1Tq1kXSth"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation results for {}\".format(dgp_types[3]))\n",
        "evals = train_setups(\n",
        "    dgp_types[3],\n",
        "    eval_outcomes,\n",
        "    features,\n",
        "    sim_samples_dict,\n",
        "    sim_samples_dict_eval,\n",
        ")\n",
        "plot(evals, dgp_types[3])"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "simulation_proxies.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
